package com.zf.agentdefine;

import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.util.Arrays;

/**
 * @author Zhou Feng
 */
public class RegisterAgent {

    /**
     * 在 启动JVM 时，通过 "-javaagent:agent-jar路径" 参数指定agent的位置，会调用这个方法：
     * <p>
     * 用于在启动时，先加载 ClassFileTransformer，在加载其他类的时候执行 ClassFileTransformer#transform() 转换类的字节码
     *
     * @param agentArgs       随同 “–javaagent”一起传入的程序参数，如果这个字符串代表了多个参数，就需要自己解析这些参数
     * @param instrumentation Instrumentation类型的对象，是JVM自动传入的，我们可以拿这个参数进行类增强等操作
     */
    public static void premain(String agentArgs, Instrumentation instrumentation) {
        System.out.printf("执行了RegisterAgent的 premain，参数是：%s\n", agentArgs);
        // 添加自定义的 ClassFileTransformer
        instrumentation.addTransformer(new TimerClassFileTransformer());
    }


    /**
     * 使用 agent-jar 修改 正在运行中的JVM进程 中的类时会调用该方法：
     * <p>
     * 用于在运行时进行类的字节码的修改，会用到 JVM Attach 机制实现
     * 步骤整体分为两步:
     * - 注册类的 ClassFileTransformer
     * - 调用 Instrumentation.retransformClasses() 方法进行类的重新加载
     *
     * @param agentArgs       使用 JVM Attach 机制加载 agent时 传入的参数
     * @param instrumentation Instrumentation类型的对象，是JVM自动传入的，我们可以拿这个参数进行类增强等操作
     */
    public static void agentmain(String agentArgs, Instrumentation instrumentation) throws UnmodifiableClassException {
        System.out.printf("执行了RegisterAgent的 agentmain，参数是：%s\n", agentArgs);

        // 获取所有被加载的类
        Class[] allLoadedClasses = instrumentation.getAllLoadedClasses();

        // 获取需要被重新加载的类
        Class[] needReloadClass = Arrays.stream(allLoadedClasses)
                .filter(item -> item.getName().equals(TimerClassFileTransformer.TRANSFORM_CLASS))
                .toArray(Class[]::new);

        // 如果没有指定的类，则直接返回
        if (needReloadClass.length == 0) {
            return;
        }

        // 添加自定义的 ClassFileTransformer，在 通过 agent 修改运行中的JVM进程中的类时，第二个参数必须是 true, 否则不会执行 ClassFileTransformer#transform()方法
        instrumentation.addTransformer(new TimerClassFileTransformer(), true);
        // 转换指定的类，并重新加载转换后的类
        instrumentation.retransformClasses(needReloadClass);

    }
}
